# 配置日志
import asyncio
import logging
import time
from typing import Optional

from protocol import MessageType, MessageProtocol, MessageBuffer, Message

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(lineno)d - %(message)s'
)
logger = logging.getLogger(__name__)


class AsyncTCPClient:
    """异步TCP客户端"""

    def __init__(self, host: str = 'localhost', port: int = 8888):
        self.host = host
        self.port = port
        self.reader: Optional[asyncio.StreamReader] = None
        self.writer: Optional[asyncio.StreamWriter] = None
        self.protocol = MessageProtocol()
        self.message_buffer = MessageBuffer()
        self.seq_counter = 0
        self.read_size = 4096
        self.connected = False
        self.receive_task: Optional[asyncio.Task] = None

        # 消息处理回调
        self.message_handlers = {
            MessageType.PING: self.handle_ping,
            MessageType.PONG: self.handle_pong,
            MessageType.DATA: self.handle_data,
            MessageType.ACK: self.handle_ack,
            MessageType.ERROR: self.handle_error,
        }

    async def handle_ping(self, message: Message):
        """处理PING消息"""
        # 自动回复PONG
        pong_msg = self.protocol.create_pong(message.seq_id)
        await self.send_message(pong_msg)
        logger.debug(f"Replied PONG to seq_id: {message.seq_id}")

    async def handle_pong(self, message: Message):
        """处理PONG消息"""
        logger.debug(f"Received PONG for seq_id: {message.seq_id}")

    async def handle_data(self, message: Message):
        """处理DATA消息"""
        try:
            data_str = message.data.decode('utf-8')
            logger.info(f"Received DATA: {data_str}")
        except UnicodeDecodeError:
            logger.warning(f"Received binary DATA: {len(message.data)} bytes")

    async def handle_ack(self, message: Message):
        """处理ACK消息"""
        logger.debug(f"Received ACK for seq_id: {message.seq_id}")

    async def handle_error(self, message: Message):
        """处理ERROR消息"""
        try:
            error_str = message.data.decode('utf-8')
            logger.error(f"Received ERROR: {error_str}")
        except UnicodeDecodeError:
            logger.error(f"Received binary ERROR: {len(message.data)} bytes")

    def get_next_seq_id(self) -> int:
        """获取下一个序列号"""
        self.seq_counter += 1
        return self.seq_counter

    async def connect(self) -> bool:
        """连接到服务器"""
        try:
            logger.info(f"Connecting to {self.host}:{self.port}...")
            self.reader, self.writer = await asyncio.open_connection(self.host, self.port)
            self.connected = True

            # 启动消息接收任务
            self.receive_task = asyncio.create_task(self.receive_messages())

            logger.info("Connected to server successfully")
            return True

        except Exception as e:
            logger.error(f"Failed to connect to server: {e}")
            return False

    async def disconnect(self):
        """断开连接"""
        logger.info("Disconnecting from server...")
        self.connected = False

        # 停止接收任务
        if self.receive_task:
            self.receive_task.cancel()
            try:
                await self.receive_task
            except asyncio.CancelledError:
                pass

        # 关闭连接
        if self.writer:
            try:
                self.writer.close()
                await self.writer.wait_closed()
            except Exception as e:
                logger.error(f"Error closing connection: {e}")

        logger.info("Disconnected from server")

    async def receive_messages(self):
        """接收消息的后台任务"""
        while self.connected:
            try:
                # 读取数据
                data = await self.reader.read(self.read_size)

                if not data:
                    logger.info("Server closed connection")
                    break

                # 解析消息
                messages = self.message_buffer.feed_data(data)

                # 处理每个完整消息
                for message in messages:
                    await self.process_message(message)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error receiving messages: {e}")
                break

        self.connected = False

    async def process_message(self, message):
        """处理接收到的消息"""
        logger.debug(f"Received message: {message.msg_type.name} seq_id: {message.seq_id}")

        # 调用对应的处理器
        handler = self.message_handlers.get(message.msg_type)
        if handler:
            try:
                await handler(message)
            except Exception as e:
                logger.error(f"Error handling message {message.msg_type.name}: {e}")
        else:
            logger.warning(f"No handler for message type: {message.msg_type.name}")

    async def send_message(self, message: Message):
        """发送消息到服务器"""
        if not self.connected or not self.writer:
            raise ConnectionError("Not connected to server")

        try:
            self.writer.write(message)
            await self.writer.drain()
        except Exception as e:
            logger.error(f"Failed to send message: {e}")
            raise

    async def send_ping(self) -> int:
        """发送PING消息"""
        seq_id = self.get_next_seq_id()
        ping_msg = self.protocol.create_ping(seq_id)
        await self.send_message(ping_msg)
        logger.debug(f"Sent PING with seq_id: {seq_id}")
        return seq_id

    async def send_data(self, data: str) -> int:
        """发送数据消息"""
        seq_id = self.get_next_seq_id()
        data_msg = self.protocol.create_data(seq_id, data.encode('utf-8'))
        await self.send_message(data_msg)
        logger.debug(f"Sent DATA with seq_id: {seq_id}")
        return seq_id


class InteractiveClient:
    """交互式客户端"""

    def __init__(self, host: str = 'localhost', port: int = 8888):
        self.client = AsyncTCPClient(host, port)
        self.running = False

    async def start(self):
        """启动交互式客户端"""
        if not await self.client.connect():
            return

        self.running = True
        logger.info("Interactive client started. Type 'help' for commands.")

        try:
            while self.running:
                try:
                    # 获取用户输入 (在实际应用中可能需要使用aioconsole)
                    command = await asyncio.get_event_loop().run_in_executor(
                        None, input, ">>> "
                    )

                    await self.process_command(command.strip())

                except KeyboardInterrupt:
                    break
                except EOFError:
                    break
                except Exception as e:
                    logger.error(f"Command error: {e}")

        finally:
            await self.client.disconnect()

    async def process_command(self, command: str):
        """处理用户命令"""
        if not command:
            return

        parts = command.split(' ', 1)
        cmd = parts[0].lower()

        if cmd == 'help':
            self.show_help()
        elif cmd == 'ping':
            await self.client.send_ping()
        elif cmd == 'send':
            if len(parts) > 1:
                await self.client.send_data(parts[1])
            else:
                print("Usage: send <message>")
        elif cmd == 'quit' or cmd == 'exit':
            self.running = False
        elif cmd == 'status':
            print(f"Connected: {self.client.connected}")
            print(f"Sequence counter: {self.client.seq_counter}")
        else:
            print(f"Unknown command: {cmd}. Type 'help' for available commands.")

    def show_help(self):
        """显示帮助信息"""
        print("""
Available commands:
  help          - Show this help message
  ping          - Send a ping message to server
  send <msg>    - Send a data message to server
  status        - Show client status
  quit/exit     - Disconnect and exit
        """)


##################################################################################################
# 简单的测试客户端

async def simple_test_client():
    """简单的测试客户端"""
    client = AsyncTCPClient()

    try:
        # 连接服务器
        if not await client.connect():
            return

        # 等待欢迎消息
        await asyncio.sleep(1)

        # 发送一些测试消息
        await client.send_data("Hello, Server!")
        await asyncio.sleep(0.5)

        await client.send_ping()
        await asyncio.sleep(0.5)

        await client.send_data("This is a test message")
        await asyncio.sleep(0.5)

        # 保持连接一段时间
        logger.info("Keeping connection alive for 10 seconds...")
        await asyncio.sleep(10)

    except Exception as e:
        logger.error(f"Test client error: {e}")
    finally:
        await client.disconnect()


# 压力测试客户端
async def stress_test_client(client_id: int, message_count: int = 100):
    """压力测试客户端"""
    client = AsyncTCPClient()

    try:
        if not await client.connect():
            return

        logger.info(f"Client {client_id} starting stress test with {message_count} messages")

        start_time = time.time()

        # 发送大量消息
        for i in range(message_count):
            await client.send_data(f"Client {client_id} message {i}")

            # 偶尔发送ping
            if i % 10 == 0:
                await client.send_ping()

            # 小延迟避免过载
            await asyncio.sleep(0.01)

        end_time = time.time()
        duration = end_time - start_time

        logger.info(f"Client {client_id} completed {message_count} messages in {duration:.2f}s")

        # 保持连接一段时间
        await asyncio.sleep(5)

    except Exception as e:
        logger.error(f"Stress test client {client_id} error: {e}")
    finally:
        await client.disconnect()


async def main():
    """主函数"""
    import sys

    if len(sys.argv) > 1:
        mode = sys.argv[1]

        if mode == 'interactive':
            # 交互式客户端
            interactive = InteractiveClient()
            await interactive.start()

        elif mode == 'test':
            # 简单测试
            await simple_test_client()

        elif mode == 'stress':
            # 压力测试
            client_count = int(sys.argv[2]) if len(sys.argv) > 2 else 5
            message_count = int(sys.argv[3]) if len(sys.argv) > 3 else 50

            logger.info(f"Starting stress test with {client_count} clients, {message_count} messages each")

            # 并发运行多个客户端
            tasks = [
                stress_test_client(i, message_count)
                for i in range(client_count)
            ]

            await asyncio.gather(*tasks)

        else:
            print("Usage: python client.py [interactive|test|stress [client_count] [message_count]]")
    else:
        # 默认运行简单测试
        await simple_test_client()


if __name__ == "__main__":
    asyncio.run(main())